import torch.utils.data as data

import os, cv2, random
import torch
import numpy as np
from ipdb import set_trace

class HeightDataset(data.Dataset):
    def __init__(self, root, txt, train = False, with_name = False, scale_size=(640, 360)):
        self.root = root
        self.scale_size = scale_size
        self.with_name = with_name
        self.txt_lines = open(txt).readlines()

    def __getitem__(self, index):
        line = self.txt_lines[index]
        fname, height = line.split(',')
        height = float(float(height) / 100.0)
        flip = False
        if np.random.rand() > 0.5:
            flip = True
        img = self.readimg(self.root, fname, flip)
        img = torch.from_numpy(img)
        rets = (img, height, )
        if self.with_name:
            rets += (fname, )
        return rets

    def readimg(self, root, fname, flip):
        img = cv2.imread(os.path.join(root, fname))
        if img.shape[0] != self.scale_size[1] or img.shape[1] != self.scale_size[0]:
            img = cv2.resize(img, self.scale_size)

        if flip:
            img = img[:,::-1,:]
        img = np.array(img, dtype=np.float32)
        img = img / 255.0 - 0.5
        img = np.transpose(img, axes=[2,0,1])
        return img

    def __len__(self):
        return len(self.txt_lines)
    
